import pandas as pd
import numpy as np
import glob
from statsmodels.tsa.statespace.sarimax import SARIMAX
import matplotlib.pyplot as plt
from datetime import datetime
import warnings

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统使用黑体
plt.rcParams['axes.unicode_minus'] = False

# 忽略警告信息
warnings.filterwarnings("ignore")

def main():
    try:
        # 1. 数据加载与预处理
        files = glob.glob('./dataset/tulufan/*.csv')  # 修改路径
        if not files:
            raise ValueError("未找到CSV文件，请检查tulufan目录")  # 修改目录名

        print("正在加载数据...")
        data_frames = []
        for file in files:
            df = pd.read_csv(file, parse_dates=['DATE'], index_col='DATE')
            data_frames.append(df)
        
        data_tulufan = pd.concat(data_frames).sort_index()  # 修改变量名
        
        # 2. 数据清洗
        print("\n数据清洗中...")
        # 2.1 处理降水量字段，将更多可能的异常值替换为np.nan
        print("  - 替换已知异常值为NaN")
        data_tulufan['PRCP'] = data_tulufan['PRCP'].replace(['999.9', '99.99', '9999', '99.9', 'nan', 'NaN'], np.nan).astype(float)
        
        # 2.2 基于统计方法过滤极端值
        print("  - 基于3σ原则检测并处理极端值")
        mean_prcp = data_tulufan['PRCP'].mean()
        std_prcp = data_tulufan['PRCP'].std()
        upper_bound = mean_prcp + 3 * std_prcp
        lower_bound = max(0, mean_prcp - 3 * std_prcp)  # 降水量不能为负
        
        # 统计极端值数量
        extreme_high_count = (data_tulufan['PRCP'] > upper_bound).sum()
        extreme_low_count = (data_tulufan['PRCP'] < lower_bound).sum()
        
        print(f"    - 检测到高于上限({upper_bound:.2f}mm)的极端值: {extreme_high_count}个")
        print(f"    - 检测到低于下限({lower_bound:.2f}mm)的极端值: {extreme_low_count}个")
        
        # 处理极端值
        data_tulufan.loc[data_tulufan['PRCP'] > upper_bound, 'PRCP'] = np.nan
        data_tulufan.loc[data_tulufan['PRCP'] < lower_bound, 'PRCP'] = np.nan
        
        # 2.3 保留必要列并重新采样
        print("  - 按天重采样并聚合数据")
        data_tulufan = data_tulufan[['PRCP']].resample('D').sum()
        
        # 2.4 缺失值处理
        print("  - 处理缺失值")
        # 统计缺失值情况
        missing_before = data_tulufan['PRCP'].isna().sum()
        print(f"    - 处理前缺失值数量: {missing_before}")
        
        # 分阶段填充缺失值
        # 第一步：将0替换为NaN（假设0表示无数据记录）
        data_tulufan['PRCP'] = data_tulufan['PRCP'].replace(0, np.nan)
        
        # 第二步：对连续缺失不超过7天的情况，使用前向填充
        data_tulufan['PRCP'] = data_tulufan['PRCP'].fillna(method='ffill', limit=7)
        
        # 第三步：对仍存在的缺失值（超过7天或其他），填充为0（假设无降水）
        data_tulufan['PRCP'] = data_tulufan['PRCP'].fillna(0)
        
        # 统计缺失值处理结果
        missing_after = data_tulufan['PRCP'].isna().sum()
        zero_after = (data_tulufan['PRCP'] == 0).sum()
        print(f"    - 处理后缺失值数量: {missing_after}")
        print(f"    - 填充为0的记录数量: {zero_after - (data_tulufan['PRCP'] == 0).sum()}")
        
        # 3. 数据检查
        print("\n数据统计信息:")
        print(f"日期范围: {data_tulufan.index.min().date()} 至 {data_tulufan.index.max().date()}")
        print(f"总记录数: {len(data_tulufan)}")
        print(f"零降水量天数: {(data_tulufan['PRCP'] == 0).sum()} ({((data_tulufan['PRCP'] == 0).sum()/len(data_tulufan))*100:.1f}%)")
        
        # 计算并打印其他统计指标
        print(f"平均降水量: {data_tulufan['PRCP'].mean():.2f} mm")
        print(f"最大降水量: {data_tulufan['PRCP'].max():.2f} mm")
        print(f"降水量中位数: {data_tulufan['PRCP'].median():.2f} mm")
        
        # 4. 可视化原始数据
        plt.figure(figsize=(12, 6))
        data_tulufan['PRCP'].plot(title='吐鲁番历史降水量数据')
        plt.xlabel('日期')
        plt.ylabel('降水量 (mm)')
        plt.grid(True)
        plt.show()

        # 5. 构建SARIMA模型
        print("\n模型训练中...")
        # 使用简化参数初始化
        model = SARIMAX(
            data_tulufan['PRCP'],
            order=(1, 1, 1),
            seasonal_order=(1, 1, 1, 30),
            enforce_stationarity=False,
            enforce_invertibility=False
        )
        
        # 模型拟合
        results = model.fit(disp=True, maxiter=100)
        print("\n模型摘要:")
        print(results.summary())

        # 6. 预测2026年清明节
        last_date = data_tulufan.index[-1]
        target_date = datetime(2026, 4, 4)
        days_ahead = (target_date - last_date).days
        
        if days_ahead <= 0:
            print("\n警告：目标日期早于现有数据的最新日期")
            return

        print(f"\n正在进行未来{days_ahead}天预测...")
        forecast = results.get_forecast(steps=days_ahead)
        forecast_mean = forecast.predicted_mean
        forecast_ci = forecast.conf_int()

        # 将预测值中的负数设为0
        forecast_mean[forecast_mean < 0] = 0
        forecast_ci[forecast_ci < 0] = 0

        # 获取清明节预测值
        qingming_forecast = forecast_mean.loc[target_date.strftime('%Y-%m-%d')]
        qingming_ci = forecast_ci.loc[target_date.strftime('%Y-%m-%d')]

        # 判断标准：降水量>0.1mm视为下雨
        is_rain = "是" if qingming_forecast > 0.1 else "否"
        
        # 7. 输出结果
        print("\n" + "="*50)
        print(f"2026年清明节（{target_date.date()}）预测结果：")
        print(f"预计降水量：{qingming_forecast:.2f} mm")
        print(f"95%置信区间：[{qingming_ci[0]:.2f} mm, {qingming_ci[1]:.2f} mm]")
        print(f"是否'雨纷纷'：{is_rain}")
        print("="*50)

        # 8. 可视化预测结果
        plt.figure(figsize=(14, 7))
        
        # 绘制最近3年数据
        recent_data = data_tulufan[data_tulufan.index >= (data_tulufan.index[-1] - pd.DateOffset(years=3))]
        plt.plot(recent_data.index, recent_data['PRCP'], label='历史数据', color='blue')
        
        # 绘制预测区间
        forecast_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=days_ahead)
        plt.plot(forecast_dates, forecast_mean, label='预测值', color='red', linestyle='--')
        plt.fill_between(forecast_dates,
                        forecast_ci.iloc[:, 0],
                        forecast_ci.iloc[:, 1],
                        color='pink', alpha=0.3, label='95%置信区间')
        
        # 标记清明节
        plt.axvline(target_date, color='green', linestyle=':', linewidth=2, label='2026清明节')
        plt.annotate(f'预测值: {qingming_forecast:.2f}mm', 
                    xy=(target_date, qingming_forecast),
                    xytext=(target_date + pd.Timedelta(days=10), qingming_forecast*1.5),
                    arrowprops=dict(facecolor='black', arrowstyle='->'))
        
        plt.title('吐鲁番降水量预测 (SARIMA模型)')
        plt.xlabel('日期')
        plt.ylabel('降水量 (mm)')
        plt.legend()
        plt.grid(True)
        plt.show()

    except Exception as e:
        print(f"\n发生错误：{str(e)}")
        print("建议排查步骤：")
        print("1. 检查数据文件是否存在及格式正确性")
        print("2. 尝试更简单的模型参数（如seasonal_order=(0,0,0,0)）")
        print("3. 安装pmdarima自动选择参数：pip install pmdarima")

if __name__ == "__main__":
    main()